#include "VulkanApplication.h"
#include "BufferManager.h"
#include "../glm/glm.hpp"
#include "../glm/gtc/matrix_transform.hpp"
#include <vector>
#include <iostream>
#include <fstream>
#include <cstring>
#include <cmath>

struct ComputePipelineData
{
    VkDescriptorPool descriptorPool;
    VkDescriptorSetLayout descriptorSetLayout;
    VkDescriptorSet descriptorSet;
    VkPipelineLayout pipelineLayout;
    VkPipeline pipeline;
    VkImage storageImage;
    VkDeviceMemory storageImageMemory;
};

VkDescriptorSetLayout createComputePipelineLayout(VulkanApplication& app, ComputePipelineData& csData)
{
    VkDescriptorSetLayoutBinding layoutBinding
    {
        0,  // binding
        VK_DESCRIPTOR_TYPE_STORAGE_IMAGE,
        1,  // descriptorCount
        VK_SHADER_STAGE_COMPUTE_BIT, nullptr
    };

    VkDescriptorSetLayoutCreateInfo layoutInfo
    {
        VK_STRUCTURE_TYPE_DESCRIPTOR_SET_LAYOUT_CREATE_INFO,
        nullptr, 0,
        1, &layoutBinding
    };

    VkDescriptorSetLayout descriptorSetLayout = nullptr;
    VkResult result = vkCreateDescriptorSetLayout(
        app.getDevice(), &layoutInfo, nullptr, &descriptorSetLayout);
    if (result != VK_SUCCESS)
    {
        printf("Failed to create description-set layout.\n");
        return nullptr;
    }

    // Create uniform desciption pool
    VkDescriptorPoolSize poolSizes[2] =
    {
        { VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER, 1 },
        { VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER, 1 }
    };

    VkDescriptorPoolCreateInfo poolInfo
    {
        VK_STRUCTURE_TYPE_DESCRIPTOR_POOL_CREATE_INFO,
        nullptr, 0,
        1, 2, poolSizes
    };

    result = vkCreateDescriptorPool(
        app.getDevice(), &poolInfo, nullptr, &csData.descriptorPool);
    if (result != VK_SUCCESS)
    {
        printf("Failed to create description pool.\n");
        return nullptr;
    }

    // Create uniform description sets
    VkDescriptorSetAllocateInfo allocInfo
    {
        VK_STRUCTURE_TYPE_DESCRIPTOR_SET_ALLOCATE_INFO,
        nullptr,
        csData.descriptorPool,
        1, &descriptorSetLayout
    };

    result = vkAllocateDescriptorSets(app.getDevice(), &allocInfo, &csData.descriptorSet);
    if (result != VK_SUCCESS)
    {
        printf("Failed to create description set.\n");
        return nullptr;
    }

    csData.descriptorSetLayout = descriptorSetLayout;
    return descriptorSetLayout;
}

bool setUpComputePipeline(VulkanApplication& app, ComputePipelineData& csData,
                          const std::string& compFile)
{
    VkShaderModule compShader = app.createShaderModule(compFile);
    VkPipelineShaderStageCreateInfo stageCreateInfo
    {
        VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO,
        nullptr, 0,
        VK_SHADER_STAGE_COMPUTE_BIT, compShader,
        "main", nullptr
    };

    VkPipelineLayoutCreateInfo pipelineLayoutInfo
    {
        VK_STRUCTURE_TYPE_PIPELINE_LAYOUT_CREATE_INFO,
        nullptr, 0,
        1, &csData.descriptorSetLayout,
        0, nullptr
    };

    VkResult result = vkCreatePipelineLayout(
        app.getDevice(), &pipelineLayoutInfo, nullptr, &csData.pipelineLayout);
    if (result != VK_SUCCESS)
    {
        printf("Failed to create pipeline layout.\n");
        return false;
    }

    VkComputePipelineCreateInfo pipelineInfo
    {
        VK_STRUCTURE_TYPE_COMPUTE_PIPELINE_CREATE_INFO,
        nullptr, 0,
        stageCreateInfo,
        csData.pipelineLayout,
        nullptr, 0
    };

    VkPipeline pipeline = nullptr;
    result = vkCreateComputePipelines(
        app.getDevice(), VK_NULL_HANDLE, 1, &pipelineInfo, nullptr, &csData.pipeline);
    if (result != VK_SUCCESS)
    {
        printf("Failed to create compute pipeline.\n");
        return false;
    }

    vkDestroyShaderModule(app.getDevice(), compShader, nullptr);
    return true;
}

VkImage createStorageImage(VulkanApplication& app, BufferManager& manager,
                           ComputePipelineData& csData, int w, int h)
{
    VkImageCreateInfo imageInfo
    {
        VK_STRUCTURE_TYPE_IMAGE_CREATE_INFO,
        nullptr, 0,
        VK_IMAGE_TYPE_2D, VK_FORMAT_R8G8B8A8_UNORM,
        { (uint32_t)w, (uint32_t)h, 1 },
        1, 1,    // mipLevels, arrayLayers
        VK_SAMPLE_COUNT_1_BIT, VK_IMAGE_TILING_OPTIMAL,
        VK_IMAGE_USAGE_STORAGE_BIT | VK_IMAGE_USAGE_SAMPLED_BIT,
        VK_SHARING_MODE_EXCLUSIVE,
        0, nullptr,    // queueFamilyIndices
        VK_IMAGE_LAYOUT_UNDEFINED
    };

    VkResult result = vkCreateImage(app.getDevice(), &imageInfo, nullptr, &csData.storageImage);
    if (result != VK_SUCCESS)
    {
        printf("Failed to create image object.\n");
        return nullptr;
    }

    VkMemoryRequirements memRequirements;
    vkGetImageMemoryRequirements(app.getDevice(), csData.storageImage, &memRequirements);
    if (!manager.createBufferMemory(csData.storageImageMemory, memRequirements,
                                    VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT))
        return nullptr;
    else
        vkBindImageMemory(app.getDevice(), csData.storageImage, csData.storageImageMemory, 0);

    // Start converting command
    VkCommandBufferAllocateInfo allocInfo
    {
        VK_STRUCTURE_TYPE_COMMAND_BUFFER_ALLOCATE_INFO,
        nullptr, app.getCommandPool(),
        VK_COMMAND_BUFFER_LEVEL_PRIMARY, 1
    };
    VkCommandBuffer tempCommand;
    vkAllocateCommandBuffers(app.getDevice(), &allocInfo, &tempCommand);

    VkCommandBufferBeginInfo beginInfo
    {
        VK_STRUCTURE_TYPE_COMMAND_BUFFER_BEGIN_INFO,
        nullptr, VK_COMMAND_BUFFER_USAGE_ONE_TIME_SUBMIT_BIT, nullptr
    };
    vkBeginCommandBuffer(tempCommand, &beginInfo);

    VkImageMemoryBarrier barrier
    {
        VK_STRUCTURE_TYPE_IMAGE_MEMORY_BARRIER,
        nullptr, 0, VK_ACCESS_SHADER_WRITE_BIT,
        VK_IMAGE_LAYOUT_UNDEFINED, VK_IMAGE_LAYOUT_GENERAL,
        VK_QUEUE_FAMILY_IGNORED, VK_QUEUE_FAMILY_IGNORED,
        csData.storageImage, { VK_IMAGE_ASPECT_COLOR_BIT, 0, 1, 0, 1 }
    };
    vkCmdPipelineBarrier(tempCommand, VK_PIPELINE_STAGE_TOP_OF_PIPE_BIT, VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT,
                         0, 0, nullptr, 0, nullptr, 1, &barrier);

    // Finish copying command and release resources
    VkQueue queue = nullptr;
    vkGetDeviceQueue(app.getDevice(), 0, 0, &queue);
    vkEndCommandBuffer(tempCommand);

    VkSubmitInfo submitInfo
    {
        VK_STRUCTURE_TYPE_SUBMIT_INFO,
        nullptr,
        0, nullptr, nullptr,  // waitSemaphores
        1, &tempCommand,
        0, nullptr    // signalSemaphore
    };
    vkQueueSubmit(queue, 1, &submitInfo, VK_NULL_HANDLE);
    vkQueueWaitIdle(queue);

    vkFreeCommandBuffers(app.getDevice(), app.getCommandPool(), 1, &tempCommand);
    return csData.storageImage;
}

int main(int argc, char** argv)
{
    VulkanApplication app("c15_compute", 1024, 768);
    if (!app.initialized())
    {
        printf("Failed to initialize VulkanApplication.\n");
        return 1;
    }

    // Create pipeline, vertices and indices
    BufferManager manager(&app);
    if (!manager.createDescriptorSet() ||
        !manager.prepare("compute.vert", "compute.frag"))
    {
        printf("Failed to create pipeline.\n");
        return 1;
    }

    // Create vertex & index buffer
    std::vector<Vertex> vertices;
    vertices.push_back({ -0.5f, -0.5f, 0.0f, 0.0f, 1.0f, 0.0f, /*uv*/ 0.0f, 0.0f });
    vertices.push_back({ 0.5f, -0.5f, 0.0f, 0.0f, 0.0f, 1.0f, /*uv*/ 1.0f, 0.0f });
    vertices.push_back({ 0.5f, 0.5f, 0.0f, 1.0f, 0.0f, 0.0f, /*uv*/ 1.0f, 1.0f });
    vertices.push_back({ -0.5f, 0.5f, 0.0f, 1.0f, 1.0f, 0.0f, /*uv*/ 0.0f, 1.0f });

    std::vector<unsigned int> indices;
    indices.push_back(0); indices.push_back(1); indices.push_back(2);
    indices.push_back(0); indices.push_back(2); indices.push_back(3);
    
    struct UniformBufferObject
    {
        glm::mat4x4 modelview;
        glm::mat4x4 projection;
    };
    UniformBufferObject ubo;

    VkBuffer vertexBuffer = manager.createVertexBuffer("vertices", vertices);
    VkBuffer indexBuffer = manager.createIndexBuffer("indices", indices);
    VkBuffer uboBuffer = manager.createUniformBuffer("ubo", sizeof(UniformBufferObject));
    void* uboData = manager.mapBufferMemory("ubo", sizeof(UniformBufferObject));

    // Create compute shader pipeline
    ComputePipelineData csData;
    createComputePipelineLayout(app, csData);
    if (!setUpComputePipeline(app, csData, "compute.comp"))
    {
        printf("Failed to create compute pipeline.\n");
        return 1;
    }

    // Create texture
    int imgW = 128, imgH = 128;
    VkImage image = createStorageImage(app, manager, csData, imgW, imgH);
    manager.createTextureSampler("image", image, VK_FORMAT_R8G8B8A8_UNORM);
    manager.updateTextureDescriptor("image", VK_IMAGE_LAYOUT_GENERAL);

    VkDescriptorImageInfo computeImageInfo
    {
        VK_NULL_HANDLE, manager.getImageView("image"), VK_IMAGE_LAYOUT_GENERAL
    };

    VkWriteDescriptorSet computeDescriptorWrite
    {
        VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET,
        nullptr,
        csData.descriptorSet, 0,  // dstSet, dstBinding
        0, 1,                     // dstArrayElement, descriptorCount
        VK_DESCRIPTOR_TYPE_STORAGE_IMAGE,
        &computeImageInfo, nullptr, nullptr
    };
    vkUpdateDescriptorSets(app.getDevice(), 1, &computeDescriptorWrite, 0, nullptr);
    
    // Main loop
    int width = 1024, height = 768;
    VkFence waitFence = app.getOrCreateFence("CommandFence");
    VkSemaphore waitNextImage = app.getOrCreateSemaphore("WaitNextImage");
    VkSemaphore waitSubmission = app.getOrCreateSemaphore("WaitSubmission");
    VkCommandBuffer commandBuffer = app.getCommandBuffer();

    float angleZ = 0.0f, aspectRatio = 0.0f;
    while (!app.shouldClose())
    {
        vkWaitForFences(app.getDevice(), 1, &waitFence, VK_TRUE, UINT64_MAX);
        vkResetFences(app.getDevice(), 1, &waitFence);

        uint32_t imageIndex = 0;
        VkResult result = vkAcquireNextImageKHR(app.getDevice(), app.getSwapchain(), UINT64_MAX,
                                                waitNextImage, VK_NULL_HANDLE, &imageIndex);
        switch (result)
        {
        case VK_SUBOPTIMAL_KHR:
        case VK_ERROR_OUT_OF_DATE_KHR:
            imageIndex = app.recreateSwapchain(width, height, waitNextImage); break;
        case VK_SUCCESS:
            break;
        default:
            printf("Next image: %d (RESULT = %d)\n", imageIndex, result);
            app.setClosed(); break;
        }

        angleZ += 0.01f; aspectRatio = (float)width / (float)height;
        glm::mat4 rotation = glm::rotate(glm::mat4(1.0f), glm::radians(angleZ), glm::vec3(0.0f, 0.0f, 1.0f));
        glm::vec3 up = glm::vec3(rotation * glm::vec4(0.0f, 1.0f, 0.0f, 1.0f));
        ubo.modelview = glm::lookAt(glm::vec3(0.0f, 0.0f, 2.0f), glm::vec3(), up);
        ubo.projection = glm::perspective(45.0f, aspectRatio, 0.1f, 100.0f);
        memcpy(uboData, &ubo, sizeof(UniformBufferObject));

        app.beginRendering(imageIndex, width, height);
        {
            VkBuffer vertexBuffers[] = { vertexBuffer };
            VkDeviceSize offsets[] = { 0 };
            VkDescriptorSet descriptorSet = manager.getDescriptorSet();
            vkCmdBindDescriptorSets(commandBuffer, VK_PIPELINE_BIND_POINT_GRAPHICS,
                                    app.getPipelineLayout(), 0, 1, &descriptorSet, 0, nullptr);

            vkCmdBindVertexBuffers(commandBuffer, 0, 1, vertexBuffers, offsets);
            vkCmdBindIndexBuffer(commandBuffer, indexBuffer, 0, VK_INDEX_TYPE_UINT32);
            vkCmdDrawIndexed(commandBuffer, (uint32_t)indices.size(), 1, 0, 0, 0);
        }
        vkCmdEndRenderPass(commandBuffer);

        {
            vkCmdBindPipeline(commandBuffer, VK_PIPELINE_BIND_POINT_COMPUTE, csData.pipeline);
            vkCmdBindDescriptorSets(commandBuffer, VK_PIPELINE_BIND_POINT_COMPUTE,
                csData.pipelineLayout, 0, 1, &csData.descriptorSet, 0, 0);
            vkCmdDispatch(commandBuffer, imgW / 8, imgH / 8, 1);

            VkImageMemoryBarrier barrier
            {
                VK_STRUCTURE_TYPE_IMAGE_MEMORY_BARRIER,
                nullptr, VK_ACCESS_SHADER_WRITE_BIT, VK_ACCESS_SHADER_READ_BIT,
                VK_IMAGE_LAYOUT_GENERAL, VK_IMAGE_LAYOUT_GENERAL,
                VK_QUEUE_FAMILY_IGNORED, VK_QUEUE_FAMILY_IGNORED,
                image, { VK_IMAGE_ASPECT_COLOR_BIT, 0, 1, 0, 1 }
            };
            vkCmdPipelineBarrier(commandBuffer, VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT,
                                 VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT,
                                 0, 0, nullptr, 0, nullptr, 1, &barrier);
        }
        vkEndCommandBuffer(commandBuffer);

        app.submitAndPresent(waitNextImage, waitSubmission, waitFence, imageIndex);
        glfwPollEvents();
    }

    vkDeviceWaitIdle(app.getDevice());
    manager.unmapBufferMemory("ubo");
    manager.cleanup();

    vkFreeMemory(app.getDevice(), csData.storageImageMemory, nullptr);
    vkDestroyImage(app.getDevice(), csData.storageImage, nullptr);
    vkDestroyDescriptorSetLayout(app.getDevice(), csData.descriptorSetLayout, nullptr);
    vkDestroyDescriptorPool(app.getDevice(), csData.descriptorPool, nullptr);
    vkDestroyPipeline(app.getDevice(), csData.pipeline, nullptr);
    vkDestroyPipelineLayout(app.getDevice(), csData.pipelineLayout, nullptr);
    return 0;
}